"""
问题起源：
我在设计基于特征图的 loss 算法时，准备使用 KL 散度
KL 散度具体计算方法此处略去，总之，需要计算 (c, h*w) @ (h*w, c) 的矩阵乘
此时，其在 h*w 维度上必须具备语义一致性
问题来了，我使用 dino 训练方法 + CNN 网络，两组特征分别是 local 特征和 global 特征
这两组特征，它本来就不应该一样呀！
如果期待空间特征具有相同的表意，至少应该确保它们源自相同的空间位置
我可以在切图处理时保留位置元素、将位置元素信息输入 loss 公式来实现这种控制，但这样做太僵硬了
一个更好的策略是，既然空间位置表意相同，那抽离的特征应该也最相似
所以干脆的，我直接提取 local feature 和 global feature 在同一 channel 上各自 h×w 维度的最相似序列
然后假言它们是匹配的不就行了吗？
当然，local 存在缺失，所以肯定不可能一一匹配上，所以我可以取前 80% 作为匹配对吧
至此，问题抽象如下：
    我有两个序列，A = [a1, a2, ..., an], B = [b1, b2, ..., bm]
    从 A、B 中各自取出 k 个元素作为子序列，记作 C、D
    使 || C - D || 最小
为简便起见，使用 L1 范数即可

相关问题：
    大佬提供了 dtw 算法作为思路，但这个算法的目的是实现两个序列对齐，与选择元素进行匹配似乎不完全相符
    有一道赛题 https://uoj.ac/problem/984 很相似，但这里 n=m=k，而我需要进行 topk 选择，似乎也不完全相符

你给我等会，这玩意是不是就是标准的匈牙利算法？
不是？等会？诶？你是认真的？诶，不是……

好啦，经过大佬、群巨巨、GPT 和本蒟蒻的协同努力，最终确定了两个可行解法：
1. 用带权二分图匹配，增设虚拟节点，设 An, Bm, k，将 Bm 扩展为 B[m + (n-k)]，增设的虚拟节点到全部 A 权值均为 0，让 A 优先匹配虚拟节点即可
2. 用带权最小费用最大流，容量限制为 k

佬佬！巨巨！
"""
import numpy as np
from scipy.optimize import linear_sum_assignment


def main():
    result = func(A=[1, 3, 7, 8, 15], B=[2, 3, 4, 6, 9], k=5)
    print(result)


def func(A, B, k):
    n = len(A)
    m = len(B)
    A = np.asarray(A)[:, None]
    B = np.asarray(B)[None, :]
    cost_matrix = abs(A - B)
    cost_matrix = np.pad(cost_matrix, [(0, 0), (0, n - k)])
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    pairs = list(zip(row_ind, col_ind))
    pairs = [(i, j) for i, j in pairs if j < m]
    return pairs


if __name__ == '__main__':
    main()
